#! /usr/bin/python

"""
The epistasis nalysis backend
"""

import time
import os
import sys
import cPickle
import tarfile
import Configuration.epistasisconfiguration as configuration
import miginterface as mig
from Gridinterface.mylogger import log
from Epicalc.jobfragmentation import create_epistasis_jobs 

########### UPDATE / STATUS ###############


def update_epistasis(epistasis_jobs, localmode=False, result_dir=None):
    """
    Update the jobs and download results if finished. Used in GUI
    """
    jobs_done = []
    if localmode:
        mig.local_mode_on()

    pkl_fil = "epi_res_pkl"#create file for epi_resutltsX.pkl
    result_dir_pkl = os.path.join(result_dir,pkl_fil)    
    if not os.path.exists(result_dir_pkl):#ender OK!
        os.mkdir(result_dir_pkl)
        
#    prog_files.extractall(path=result_dir_pkl)??????????????????
    
    for j in epistasis_jobs:
        if j["status"] == "FINISHED": # dont update if the un-updated job has already finished  
            continue
        
        info = mig.job_info(j["id"])
        j["status"] = info["STATUS"]
        
        if j["status"] == "FINISHED": # check if the newly updated job has finished 
            j["finished"] = time.asctime()
            jobs_done.append(j)
            
            for f in j["output_files"]:
                mig.get_file(f, result_dir)            
                result_archive =  os.path.join(result_dir, f)#epifiles1.tar.gz
                #Files are extracted with this call
                extract_output(result_archive, result_dir,result_dir_pkl,j)
                os.remove(result_archive)#The tar-file is removed
                
#    if len(epistasis_jobs) == len(jobs_done):
 #       log(logfile, 'All jobs completed', debug)
    
    return epistasis_jobs

def extract_output(tar_file_path, dest_dir,result_dir_pkl,j):
    """
    Unpack the tar file
    """
    prog_files = tarfile.open(tar_file_path, "r")
#    prog_files.extractall(path=dest_dir)
    
    #Extract epi_resultsX.pkl to epi_res_pkl
    epires_pkl = "epi_results"+j["class"][0]+".pkl"
    prog_files.extract(epires_pkl,path=result_dir_pkl)        

    #Extract al eccept scripts
    for i in prog_files.getnames():
        if i[-3:] == "pyc" or i[-3:] == ".py":# or i[:3] == "epi":
            continue
        else:    
            prog_files.extract(i,path=dest_dir)     
               
    prog_files.close()
#    print "Extracted %s to %s ." % (tar_file_path, dest_dir)
 

def create_epi_mig_job(job, localmode=False):
    """
    Create a mig job
    """
    if localmode:
        mig.local_mode_on()
    job["id"]= mig.create_job(exec_commands=job["commands"], input_files=job["input_files"],                                     
                                  output_files=job["output_files"],
                                  resource_specifications=job["resource_specs"])
    job["started"] = time.asctime()
    return job

# ######### START EPISTASIS ############


def start_epistasis_from_UI(user_args, progress_callback):
    """Start the epistasis analysis."""

    selection_variable = user_args["class_type"]
    classes = user_args["classes"] #selection_variable_values = configuration.default_variable_values,
    genotype_names = user_args["snp_gene_names"]
    genotype_matrixes = user_args["genotype_files"] #genelist=configuration.default_gene_list,
    phenotype_names = user_args["phenotype_names"]
    phenotype_matrixes = user_args["phenotype_files"] #traitlist=configuration.default_trait_list,
    name = user_args["analysis_name"] # user specified name of the analysis
    local_mode = user_args["local_mode"] #local_mode=False,
    status_file = user_args["status_file"] # pkl file to store status information
    result_dir = user_args["result_dir"] # result directory
    bonferoni_corr = user_args["bonferoni_correction"]
    significance_level = user_args["significance_level"]

    job_size = 1
    
    if os.path.exists(result_dir):
        raise Exception("result dir already exists : "+result_dir)

    # the result dir is location where the tar archives are unpacked
    os.mkdir(result_dir)

    if os.path.exists(status_file):
        raise Exception("status file exists: "+status_file)
    
    # create the epistastis jobs
    epi_jobs = create_epistasis_jobs(
        job_size,
        genotype_names, 
        genotype_matrixes,
        phenotype_names,
        phenotype_matrixes,
        selection_variable,
        classes,
        significance_level=significance_level,
        bonf_corr=bonferoni_corr,
        analysis_name=name,
        run_local=local_mode
        )
    
    if local_mode:
        add_alias()
        
    migjobs = []   
    job_nr = 0
    # create and submit the epistasis jobs
    for j in epi_jobs:
        job_nr +=1
        
        migjobs.append(create_epi_mig_job(j, local_mode))
        progress_callback("Creating epistasis job (%i of %i)"% (job_nr, len(epi_jobs)), float(job_nr)/float(len(epi_jobs))*80) # 80% is max

    # create a status file     
    analysis_status = {}
    analysis_status["local"] = local_mode
    analysis_status["jobs"] = migjobs
    analysis_status["start_time"] = time.asctime()
    analysis_status["result_dir"] = result_dir
    
    f = open(status_file, "w")
    cPickle.dump(analysis_status, f)
    f.close()
    
    #progress_callback("Startup complete.", 100)


########## STOP /CANCEL ##############

def stop_epistasis(jobs):
    """Stop the epistasis procedure."""
    #mig_session.cancel_jobs(jobs)
    for j in jobs:
        mig.cancel_job(j["id"])
    

###### PRINT ###########

def print_jobs(self, jobs):
    """Print jobs."""
    for i in range(len(jobs)):
        print 'job ' + str(i) + ' : ' + str(jobs[i])

def print_status(self, jobs):
    """Print job status."""
    full_str = []
    for j in jobs:
        status_str = 'Job : ' + j['id'] + '\t' + j['status'
                ]['STATUS']
        print status_str
        full_str.append(status_str)
    return full_str

# ### CLEAN UP ########

def clean_up_epistasis(jobs):
    """Delete files used in the epistasis procedure that are no longer needed."""
    mig_session.clean_up(jobs)

def add_alias():
    os.putenv("PYTHON", "/usr/bin/python")
#    os.putenv("R_HOME", "/usr/lib/R")

###### MAIN ###############

# Arguments are entered in the order: selectionvariableindex jobsize

#local = False
debug = False

if '-local' in sys.argv or '-l' in sys.argv:
        local = True
        # set up the needed mig job runtime environment
        os.putenv("PYTHON", "/usr/bin/python")
        os.putenv("R_HOME", "/usr/lib/R")
if '-debug' in sys.argv or '-d' in sys.argv:
        debug = True
#logfile = configuration.output_dir+configuration.logfile_name
#mig_session = migsession.Migsession(configuration.output_dir, logfile, local, debug)

if __name__ == '__main__':
    if '-u' in sys.argv :
        i = sys.argv.index("-u")
        pkl_file = open(sys.argv[i+1], "r+")
        jobs = cPickle.load(pkl_file)
 #       get_epistasis_status(jobs)
#        print_epistasis_status(jobs)
        cPickle.dump(jobs, pkl_file)
        pkl_file.close()
    
    elif '-c' in sys.argv :
        i = sys.argv.index("-c")
        pkl_file = open(sys.argv[i+1], "r+")
        jobs = cPickle.load(pkl_file)
        #get_epistasis_status(jobs)
        stop_epistasis(jobs)
 #       print_epistasis_status(jobs)
        cPickle.dump(jobs, pkl_file)
        pkl_file.close()
    
    else:
        start_epistasis(local_mode=local)
